/**
  ETFRoot project, Anyang Normal University && IMP-CAS
  \class ETFPID
  \brief class for particle identification. This class takes tracks on both sides
  of the dipole magnet to extract particle A/Z, and particle velocity beta.
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2021-11-10
  \date 2021-12-15 last modified
  \attention
  changelog
  <table>
  <tr>  <th>Date         <th>Author      <th>Description                   </tr>
  <tr>  <td>2021-11-10   <td>Asia Sun    <td>file created                  </tr>
  </table>

  \copyright Copyright (c) 2021 Anyang Normal U. && IMP-CAS with LGPLv3 LICENSE
*/

#include <iostream>
#include <TGraph.h>
#include <TLegend.h>
#include <TClonesArray.h>
#include "FairRootManager.h"
#include "ETFPID.h"
#include "ETFTrackManager.h"
#include "ETFDetectorManager.h"
#include "ETFPreTaTracker.h"
#include "ETFPostTaTracker.h"
#include "ETFPostMagTracker.h"
#include "ETFGlobalTracker.h"
#include "ETFTrackTa4.h"
#include "ETFTrackTa.h"
#include "ETFTrack.h"
#include "ETF3DTrack.h"
#include "ETFTOFWall.h"
#include "ETFTstop.h"
#include "ETFTstart.h"
#include "ETFMath.h"
#include "ETFMsg.h"
#include "ETFPlaHit.h"
#include "ETFParaManager.h"
#include "ETFYaml.h"
#include "ETFVMEEvent.h"

using std::cout;
using std::endl;

ClassImp(ETFPID);

double ETFPID::kB = -9999.;
static constexpr double L = 25881.; // distance from Tstart to Tstop in mm
static constexpr double c0 = ETFMath::c0(); // in mm/ns
static constexpr double u0 = ETFMath::uMeV(); // in MeV
#define sum2 ETFMath::Sum2

ETFPID::ETFPID() : TNamed("pid", "PID"), fGCurve(nullptr), fGLegend(nullptr){
  Initialize();
  ftof2_ofst = -9999.;
} // end ctor

ETFPID::~ETFPID(){}

void ETFPID::Initialize(){
  fHasFlied = false;
  fdsca11 = -9999; fdsca4 = -9999;
  ftof1 = -9999.; fBG0 = fGamma0 = fBeta0 = -9999.;
  ftof2 = -9999.; fBG  = fGamma  = fBeta  = -9999.;
  fdE0 = fdE1 = -9999.;

  fAoZ = -9999.; fBrho = -9999.; fPoZ = -9999.;
  // x2: abscissa of Mag exit point from real trk and calculated arc,
  // as a PID estimator; 0-1: real-calcu
  fX2Arr[0] = -9999.; fX2Arr[1] = -9999.;
  fZO = -9999.; fXO = -9999.; fRho = -9999.;
  fTrkLenT = -9999.;

  fCospl = fdCospl = fPoZpl = fBGC = fdBGC = -9999.;
} // end member function Initialize

// propogate the particle in the dipole magnet to meet the tracks given by the DCs
// currently and usually a uniform magnetic field is assumed, as it is both
// appropriate and efficient
// tof2: from the target to TOFWall
// p0[0-1]: {k, b} from postTaTrk
// p1[0-1]: {k, b} from postMagTrk
bool ETFPID::Fly(const double *p0, const double *p1, double tof2){
  if(fHasFlied) return true;

  static ETFDetectorManager *m = ETFDetectorManager::Instance();
  static const double zTstop = m->GetDetector("tstop")->GetZ();
  static const double zTarget = m->GetTargetZ();
  // exit point of magField
  static const double zi = ETFMath::kzMagIn, ze = ETFMath::kzMagOut;

  // rough estimate particle beta to decide whether it is appropriate to proceed PID //
  const double xe = p1[0]*ze+p1[1]; // x of the point exiting the magField
  double pe[3]{}; // the position of TOFWall hit
  dynamic_cast<ETFTOFWall *>(m->GetDetector("tofw"))->GetHitPosition(p1, pe);
  // do the particle propogation in magField assuming a uniform magField //
  fRho = ETFMath::rho(p0[0], p0[1], p1[0], p1[1], fX2Arr, &fZO, &fXO); // in mm
  fTrkLenT = (zi-zTstop)*sqrt(1.+p0[0]*p0[0]) + fabs(fRho*(atan(p1[0])-atan(p0[0])))
    + sqrt(sum2(ze-pe[2], xe-pe[0]));
  fBeta = fTrkLenT / (c0*tof2);
  if(fBeta >= 1. || fBeta < 0.){ // abnormal beta
    // fBeta = -1.;
    return false;
  } // end if
  fBG = ETFMath::BetaGamma(fBeta); fGamma = fBG/fBeta;
  // 0.321840605 = e0/(u0*c0*1E6) SI unit
  static constexpr double CC = 0.32184060506023993;
  fBrho = kB*fRho*1.e-3; // in T.m; r: 1e-3 factor for mm->m
  fPoZ = fBrho*c0; // in MeV/c
  fAoZ = fBrho*CC/fBG;
  fHasFlied = true;

  return true;
} // end member function Fly

// the same fly, but collecting information using singleton classes in ETFAna
// istof: whether we only call fly to obtain tof
bool ETFPID::Fly(bool istof){
  if(fHasFlied || (istof && -9999. != ftof2)) return true; // has called

  static ETFDetectorManager *dm = ETFDetectorManager::Instance();
  static ETFTrackManager *tm = ETFTrackManager::Instance();
  static ETFParaManager *pm = ETFParaManager::Instance();
  if(-9999. == ftof2_ofst) ftof2_ofst = (*pm->GetDelayPara())["tof2_ofst"].as<double>();
  tm->Map(); // in case that tracking has not been done

  static ETFGlobalTracker *gt = tm->GetGlobalTracker();
  static ETFPostMagTracker *pmt = tm->GetPostMagTracker();
  // even tof2 cannot be calculated without (only 1) trk-postMag
  if(1 != pmt->GetNofTracks()){ fHasFlied = true; return false; }
  // then we can calculate tof2 //
  ETFTrack *trkpm = dynamic_cast<ETFTrack *>(pmt->GetTrack(0));
  const int stripId = trkpm->GetFiredStripId();
  ETFTstart *tstart = dynamic_cast<ETFTstart *>(dm->GetDetector("tstart"));
  ETFTstop *tstop = dynamic_cast<ETFTstop *>(dm->GetDetector("tstop"));
  ETFTOFWall *tofw = dynamic_cast<ETFTOFWall *>(dm->GetDetector("tofw"));
  // not properly fired -- not even ready to give TOF information
  if(4 != tstop->GetFiredStatus() || !tstart->GetFiredStatus()) return false;
  ftof1 = tstop->GetTime() - tstart->GetTime();
  ftof2 = tofw->GetStrip(stripId)->GetTime() - tstop->GetTime() + ftof2_ofst;
  fBeta0 = L / (ftof1*c0);
  if(fBeta0 >= 1. || fBeta0 < 0.){ // abnormal beta0
    // fBeta0 = -1.;
    return false;
  } // end if
  fBG0 = ETFMath::BetaGamma(fBeta0); fGamma0 = fBG0/fBeta0;
  if(istof) return true;

  // PID condition: only 1 track in globalTracker and only 1 track in postMagTracker
  if(1 != gt->GetNofTracks()){ fHasFlied = true; return false; }
  ETFTrackTa *trkbt = dynamic_cast<ETFTrackTa4 *>(gt->GetTrack(0))->GetPreTaTrack();
  ETFTrackTa *trkpt = dynamic_cast<ETFTrackTa4 *>(gt->GetTrack(0))->GetPostTaTrack();
  const double p0[2] = {trkpt->GetK(), trkpt->GetB()};
  const double p1[2] = {trkpm->GetK(), trkpm->GetB()};

  const bool fly = Fly(p0, p1, ftof2);
  costh(); // calculate the parallel momentum of the residue

  return fly;
} // end member function Fly

void ETFPID::SetPileUp(ETFVMEEvent *vme){
  Setdsca11(vme->GetDSca()[11]);
  Setdsca4(vme->GetDSca()[4]);
} // end member function SetPileUp

// cos(theta_c-theta_p)
bool ETFPID::costh(){
  static ETFTrackManager *tm = ETFTrackManager::Instance();
  static ETFPreTaTracker *preTaTracker = tm->GetPreTaTracker();
  static ETFPostTaTracker *postTaTracker = tm->GetPostTaTracker();

  // validity check
  if(1 != preTaTracker->GetNof3DTracks()) return false;
  if(1 != postTaTracker->GetNof3DTracks()) return false;
  if(-9999. == fPoZ) return false;

  ETF3DTrack *trk0 = dynamic_cast<ETF3DTrack *>(preTaTracker->Get3DTrack(0));
  ETF3DTrack *trk1 = dynamic_cast<ETF3DTrack *>(postTaTracker->Get3DTrack(0));
  const double k01 = trk0->GetK1(), k02 = trk0->GetK2();
  const double k11 = trk1->GetK1(), k12 = trk1->GetK2();
  const double n0[3] = {1, k01, k02}, n1[3] = {1, k11, k12};
  const double n0m = ETFMath::norm(n0), n1m = ETFMath::norm(n1); // |n0| and |n1|
  // pl: parallel
  fCospl = ETFMath::innerProd(n0,n1) / (n0m*n1m); // n0*n1/(|n0|*|n1|)=cos<n0,n1>
  // calculate dcospl to see its range of magnitude //
  // slope resolution (sigma) of core (C) track and projectile (P) track around TA
  static const double DKC = 0.0001, DKP = 0.0001; // preTa: 3.07201e-03
  const double t1 = DKP*n1m*n1m, t2 = DKC*n0m*n0m;
  fdCospl = sqrt(sum2(
    // dcospl/dk01   dcospl/dk02
    // dcospl/dk11   dcospl/dk12
    t1*((1.+k02*k02)*k11-k01*(1.+k02*k12)),t1*((1.+k01*k01)*k12-k02*(1.+k01*k11)),
    t2*((1.+k12*k12)*k01-k11*(1.+k12*k02)),t2*((1.+k11*k11)*k02-k12*(1.+k11*k01))));
  fdCospl /= pow(n0m*n1m,3);
  // calculate dpozpl to see its range of magnitude //
  static const double DBG = 2.00767e-3; // sigma(beta*gamma)
  /// fdBGC*aoz=fdpozpl, yet more accurate (aoz resolution excluded)
  /// by using real aoz after identifying the nuclide
  fBGC = fBG*fCospl*u0; // fBGC*aoz gives more accurate fPoZpl with real aoz
  fdBGC = fBGC*sqrt(sum2(DBG/fBG, fdCospl/fCospl));
  static const double DRHOREL = 0.01; // sigma(rho)/rho ~ 0.01 for ETF
  fPoZpl = fPoZ * fCospl; // projection to the incident direction
  fdPoZpl = fPoZpl*sqrt(sum2(fdCospl/fCospl, DRHOREL));
  // cout << endl << "DBG/fBG: " << DBG/fBG; // DEBUG
  // cout << " fdCospl/fCospl: " << fdCospl/fCospl << endl; // DEBUG
  // cout << "(DBG/fBG)/(fdCospl/fCospl): " << (DBG/fBG)/(fdCospl/fCospl) << endl; // DEBUG
  // cout << "fdBGC/fBGC: " << fdBGC/fBGC << endl; // DEBUG
  // cout << "fdPoZpl/fPoZpl: " << fdPoZpl/fPoZpl << endl; // DEBUG
  // getchar(); // DEBUG

  return true;
} // end member function cost


void ETFPID::Draw(){
  ConstructGraph();
  fGCurve->Draw("plsame");
  if(-9999. != fAoZ) fGLegend->Draw();
} // end member function Draw

void ETFPID::ConstructGraph(){
  if(!fHasFlied) return; // nothing to draw
  // draw the arc
  if(!fGCurve){
    fGCurve = new TGraph();
    fGCurve->SetMarkerStyle(1); fGCurve->SetMarkerSize(3);
    fGCurve->SetMarkerColor(6); fGCurve->SetLineColor(6);
  } // end if
  if(!fGLegend){
    // mark out the aoz value of the event
    fGLegend = new TLegend(0.4, 0.8, 0.6, 0.89); fGLegend->SetBorderSize(0);
    fGLegend->AddEntry(fGCurve, "", "lp");
  } // end if

  while(fGCurve->GetN()) fGCurve->RemovePoint(0); // clear the graph
  const int n = 400; // n points to draw the curve
  for(int i = 0; i < n; i++){
    const double ai = (1.-2.*i/n)*ETFMath::Pi();
    const double zi = fZO + fRho*cos(ai);
    const double xi = fXO + fRho*sin(ai);
    fGCurve->SetPoint(i, zi, xi);
  } // end for over i
  // label the aoz value in the graph
  if(-9999. == fAoZ) return;
  char name[64]; sprintf(name, "aoz_%1.3f\n", fAoZ);
  fGLegend->SetEntryLabel(name);
} // end member function ConstructGraph

// x2arr[0] - x2arr[1]
double ETFPID::dx2() const{
  if(-9999. == fX2Arr[0] || -9999. == fX2Arr[1]) return -9999.;
  return fX2Arr[0] - fX2Arr[1];
} // end member function dx2

void ETFPID::Print(){
  cout << "\033[32;1m============= PID PRINT =============\033[0m" << endl;
  cout << "aoz: " << fAoZ << " beta: " << fBeta;
  cout << " poz: " << fPoZ << " fBrho: " << fBrho << endl;
  cout << "total track length(tof2): " << fTrkLenT << endl;
  cout << "x2: " << fX2Arr[0] << " " << fX2Arr[1] << " dx2: " << fX2Arr[0] - fX2Arr[1] << endl;
  cout << "tof1: " << ftof1 << "  tof2: " << ftof2 << endl;
  cout << "beta0: " << fBeta0 << "  rho: " << fRho << endl;
  cout << "dE0: " << fdE0 << " dE1: " << fdE1 << endl;
  cout << "dsca4: " << fdsca4 << " dsca11: " << fdsca11 << endl;
  cout << "------ Parallel Momentum --------" << endl;
  cout << "cospl: " << fCospl << "("<< fdCospl << ")" << endl;
  cout << "BGC: " << fBGC << "("<< fdBGC << ")" << endl;
  cout << "pozpl: " << fPoZpl << "("<< fdPoZpl << ")" << endl;
} // end member function Print
